import os
import json
import random

dataset_path = '/home/ruideng/workspace/Ali_Product'
dst_path = '/home/ruideng/workspace/LSPR'

def json2listfile(src_path=dataset_path, dst_path=dst_path):
    file_name = ['train', 'val']
    for file in file_name:
        src_file = os.path.join(src_path, file+'.json')
        contents = json.load(open(src_file))

        dst_fp = open(os.path.join(dst_path, file+'.txt'), 'w')
        contents = contents['images']

        for line in contents:
            dst_fp.writelines(os.path.join(dataset_path, file, line['class_id'], line['image_id'])
            + ' ' + line['class_id'] + '\n')
        
        dst_fp.close()

def json2listfile_subset(src_path=dataset_path, dst_path=dst_path, classes=(50030, 500)):
    file_name = ['train', 'val']

    cls_choice = list(range(classes[0]))
    random.shuffle(cls_choice)
    cls_choice = cls_choice[0:classes[1]]

    for file in file_name:
        src_file = os.path.join(src_path, file+'.json')
        contents = json.load(open(src_file))

        dst_fp = open(os.path.join(dst_path, file+'_tiny{}.txt'.format(classes[1])), 'w')
        contents = contents['images']

        for line in contents:
            if int(line['class_id']) in cls_choice:
                dst_fp.writelines(os.path.join(dataset_path, file, line['class_id'], line['image_id'])
                + ' ' + line['class_id'] + '\n')
        
        dst_fp.close()

if __name__ == '__main__':
    json2listfile()
